"""
WebSocket实时流服务
"""

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, Query, HTTPException, status
from typing import Dict, List, Any, Optional
import json
import asyncio
import uuid
from datetime import datetime
import logging

from ...core.exceptions import ValidationError, BusinessLogicError, with_error_handling
from ...core.security import get_current_user, rate_limit
from ...core.logging import audit_logger
from ...services.audio import StreamManager, StreamConfig, StreamStatus
from ...services.tts import tts_service


# 创建路由
router = APIRouter()

# 全局WebSocket连接管理
active_connections: Dict[str, Dict[str, Any]] = {}


class WebSocketConnection:
    """WebSocket连接管理类"""
    
    def __init__(self, websocket: WebSocket, connection_id: str, user_id: str):
        self.websocket = websocket
        self.connection_id = connection_id
        self.user_id = user_id
        self.connected_at = datetime.utcnow()
        self.last_activity = datetime.utcnow()
        self.stream_sessions: Dict[str, StreamManager] = {}
        self.is_authenticated = True
    
    async def send_message(self, message: Dict[str, Any]):
        """发送消息到客户端"""
        try:
            await self.websocket.send_text(json.dumps(message, ensure_ascii=False))
            self.last_activity = datetime.utcnow()
        except Exception as e:
            print(f"发送WebSocket消息失败: {e}")
    
    async def send_error(self, error_message: str, error_code: str = None):
        """发送错误消息"""
        error_data = {
            "type": "error",
            "message": error_message,
            "timestamp": datetime.utcnow().isoformat()
        }
        if error_code:
            error_data["error_code"] = error_code
        
        await self.send_message(error_data)
    
    async def send_stream_update(self, stream_id: str, status: str, data: Dict[str, Any]):
        """发送流更新消息"""
        update_data = {
            "type": "stream_update",
            "stream_id": stream_id,
            "status": status,
            "data": data,
            "timestamp": datetime.utcnow().isoformat()
        }
        await self.send_message(update_data)
    
    async def heartbeat(self):
        """心跳检测"""
        heartbeat_data = {
            "type": "heartbeat",
            "timestamp": datetime.utcnow().isoformat()
        }
        await self.send_message(heartbeat_data)
    
    def update_activity(self):
        """更新活动时间"""
        self.last_activity = datetime.utcnow()


class WebSocketManager:
    """WebSocket管理器"""
    
    def __init__(self):
        self.connections: Dict[str, WebSocketConnection] = {}
        self.user_connections: Dict[str, List[str]] = {}  # user_id -> [connection_ids]
    
    async def connect(self, websocket: WebSocket, user_id: str) -> str:
        """建立WebSocket连接"""
        connection_id = str(uuid.uuid4())
        
        await websocket.accept()
        
        # 创建连接对象
        connection = WebSocketConnection(websocket, connection_id, user_id)
        
        # 保存连接
        self.connections[connection_id] = connection
        
        # 更新用户连接列表
        if user_id not in self.user_connections:
            self.user_connections[user_id] = []
        self.user_connections[user_id].append(connection_id)
        
        # 发送连接确认
        await connection.send_message({
            "type": "connected",
            "connection_id": connection_id,
            "user_id": user_id,
            "timestamp": datetime.utcnow().isoformat()
        })
        
        # 记录审计日志
        audit_logger.log_user_action(
            user_id=user_id,
            action="websocket_connect",
            resource="websocket",
            details={"connection_id": connection_id}
        )
        
        return connection_id
    
    def disconnect(self, connection_id: str):
        """断开WebSocket连接"""
        if connection_id not in self.connections:
            return
        
        connection = self.connections[connection_id]
        user_id = connection.user_id
        
        # 清理流会话
        for stream_id, stream_manager in connection.stream_sessions.items():
            asyncio.create_task(stream_manager.cleanup())
        
        # 从连接列表中移除
        del self.connections[connection_id]
        
        # 从用户连接列表中移除
        if user_id in self.user_connections:
            if connection_id in self.user_connections[user_id]:
                self.user_connections[user_id].remove(connection_id)
            
            # 如果用户没有其他连接，移除用户记录
            if not self.user_connections[user_id]:
                del self.user_connections[user_id]
        
        # 记录审计日志
        audit_logger.log_user_action(
            user_id=user_id,
            action="websocket_disconnect",
            resource="websocket",
            details={"connection_id": connection_id}
        )
    
    def get_connection(self, connection_id: str) -> Optional[WebSocketConnection]:
        """获取连接对象"""
        return self.connections.get(connection_id)
    
    def get_user_connections(self, user_id: str) -> List[WebSocketConnection]:
        """获取用户的所有连接"""
        if user_id not in self.user_connections:
            return []
        
        connections = []
        for conn_id in self.user_connections[user_id]:
            if conn_id in self.connections:
                connections.append(self.connections[conn_id])
        
        return connections
    
    async def broadcast_to_user(self, user_id: str, message: Dict[str, Any]):
        """向用户的所有连接广播消息"""
        connections = self.get_user_connections(user_id)
        
        tasks = []
        for connection in connections:
            tasks.append(connection.send_message(message))
        
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)
    
    async def cleanup_inactive_connections(self, timeout_seconds: int = 300):
        """清理不活跃的连接"""
        now = datetime.utcnow()
        inactive_connections = []
        
        for conn_id, connection in self.connections.items():
            if (now - connection.last_activity).total_seconds() > timeout_seconds:
                inactive_connections.append(conn_id)
        
        for conn_id in inactive_connections:
            self.disconnect(conn_id)
        
        return len(inactive_connections)


# 创建全局WebSocket管理器
websocket_manager = WebSocketManager()


# WebSocket端点
@router.websocket("/stream")
@with_error_handling("websocket_stream")
async def websocket_stream_endpoint(
    websocket: WebSocket,
    token: Optional[str] = Query(None),
    user_id: Optional[str] = Query(None)
):
    """WebSocket流传输端点"""
    try:
        # 验证用户身份
        current_user = None
        if token:
            try:
                from ...core.security import token_manager
                payload = token_manager.verify_token(token)
                current_user = payload
                user_id = payload.get("sub", user_id)
            except Exception:
                await websocket.close(code=4001, reason="Invalid token")
                return
        
        if not user_id:
            await websocket.close(code=4002, reason="User ID required")
            return
        
        # 建立连接
        connection_id = await websocket_manager.connect(websocket, user_id)
        connection = websocket_manager.get_connection(connection_id)
        
        # 启动心跳检测
        heartbeat_task = asyncio.create_task(heartbeat_loop(connection))
        
        try:
            # 处理客户端消息
            while True:
                message = await websocket.receive_text()
                connection.update_activity()
                
                try:
                    data = json.loads(message)
                    await handle_client_message(connection, data, current_user)
                except json.JSONDecodeError:
                    await connection.send_error("Invalid JSON message")
                except Exception as e:
                    await connection.send_error(f"Message processing error: {str(e)}")
        
        except WebSocketDisconnect:
            pass
        finally:
            # 清理
            heartbeat_task.cancel()
            websocket_manager.disconnect(connection_id)
            
    except Exception as e:
        print(f"WebSocket连接错误: {e}")


async def handle_client_message(
    connection: WebSocketConnection,
    data: Dict[str, Any],
    user: Optional[Dict[str, Any]]
):
    """处理客户端消息"""
    message_type = data.get("type")
    
    if message_type == "start_stream":
        await handle_start_stream(connection, data, user)
    elif message_type == "stop_stream":
        await handle_stop_stream(connection, data)
    elif message_type == "pause_stream":
        await handle_pause_stream(connection, data)
    elif message_type == "resume_stream":
        await handle_resume_stream(connection, data)
    elif message_type == "get_stream_status":
        await handle_get_stream_status(connection, data)
    elif message_type == "ping":
        await handle_ping(connection, data)
    else:
        await connection.send_error(f"Unknown message type: {message_type}", "UNKNOWN_MESSAGE_TYPE")


async def handle_start_stream(
    connection: WebSocketConnection,
    data: Dict[str, Any],
    user: Optional[Dict[str, Any]]
):
    """处理开始流消息"""
    try:
        stream_config = data.get("config", {})
        text = data.get("text")
        voice_id = data.get("voice_id")
        voice_preset = data.get("voice_preset")
        
        if not text:
            await connection.send_error("Text is required for streaming", "MISSING_TEXT")
            return
        
        # 创建流配置
        config_data = {
            "buffer_size": stream_config.get("buffer_size", 8192),
            "max_buffers": stream_config.get("max_buffers", 10),
            "sample_rate": stream_config.get("sample_rate", 22050),
            "channels": stream_config.get("channels", 1),
            "format": stream_config.get("format", "wav"),
            "chunk_duration": stream_config.get("chunk_duration", 0.1),
            "preload_chunks": stream_config.get("preload_chunks", 3),
            "auto_reconnect": stream_config.get("auto_reconnect", True),
            "reconnect_delay": stream_config.get("reconnect_delay", 1.0)
        }
        
        stream_config_obj = StreamConfig(**config_data)
        
        # 创建流管理器
        stream_id = str(uuid.uuid4())
        stream_manager = StreamManager(stream_config_obj)
        
        # 设置回调
        await setup_stream_callbacks(connection, stream_manager, stream_id)
        
        # 创建音频源
        async def audio_source():
            # 获取语音配置
            voice_profile = None
            if voice_id:
                voices = await tts_service.get_available_voices()
                voice_profile = next((v for v in voices if v.id == voice_id), None)
            
            # 应用语音预设
            if voice_preset:
                await tts_service.set_voice_preset(voice_preset)
            
            # 执行TTS
            response = await tts_service.synthesize(
                text=text,
                voice_profile=voice_profile,
                voice_settings=tts_service.voice_controller.settings,
                output_format=stream_config_obj.format
            )
            
            if response.success and response.audio_data:
                return response.audio_data
            else:
                raise BusinessLogicError(f"TTS failed: {response.error_message}")
        
        # 保存流管理器
        connection.stream_sessions[stream_id] = stream_manager
        
        # 开始流传输
        success = await stream_manager.start_streaming(audio_source)
        
        if success:
            await connection.send_message({
                "type": "stream_started",
                "stream_id": stream_id,
                "config": config_data,
                "timestamp": datetime.utcnow().isoformat()
            })
            
            # 记录审计日志
            audit_logger.log_user_action(
                user_id=connection.user_id,
                action="start_stream",
                resource="websocket",
                details={
                    "stream_id": stream_id,
                    "text_length": len(text),
                    "voice_id": voice_id
                }
            )
        else:
            del connection.stream_sessions[stream_id]
            await connection.send_error("Failed to start stream", "STREAM_START_FAILED")
            
    except Exception as e:
        await connection.send_error(f"Start stream error: {str(e)}", "START_STREAM_ERROR")


async def handle_stop_stream(connection: WebSocketConnection, data: Dict[str, Any]):
    """处理停止流消息"""
    stream_id = data.get("stream_id")
    
    if not stream_id:
        await connection.send_error("Stream ID is required", "MISSING_STREAM_ID")
        return
    
    if stream_id not in connection.stream_sessions:
        await connection.send_error(f"Stream {stream_id} not found", "STREAM_NOT_FOUND")
        return
    
    stream_manager = connection.stream_sessions[stream_id]
    await stream_manager.stop_streaming()
    await stream_manager.cleanup()
    
    del connection.stream_sessions[stream_id]
    
    await connection.send_message({
        "type": "stream_stopped",
        "stream_id": stream_id,
        "timestamp": datetime.utcnow().isoformat()
    })


async def handle_pause_stream(connection: WebSocketConnection, data: Dict[str, Any]):
    """处理暂停流消息"""
    stream_id = data.get("stream_id")
    
    if not stream_id or stream_id not in connection.stream_sessions:
        await connection.send_error("Invalid stream ID", "INVALID_STREAM_ID")
        return
    
    stream_manager = connection.stream_sessions[stream_id]
    await stream_manager.pause_streaming()
    
    await connection.send_message({
        "type": "stream_paused",
        "stream_id": stream_id,
        "timestamp": datetime.utcnow().isoformat()
    })


async def handle_resume_stream(connection: WebSocketConnection, data: Dict[str, Any]):
    """处理恢复流消息"""
    stream_id = data.get("stream_id")
    
    if not stream_id or stream_id not in connection.stream_sessions:
        await connection.send_error("Invalid stream ID", "INVALID_STREAM_ID")
        return
    
    stream_manager = connection.stream_sessions[stream_id]
    await stream_manager.resume_streaming()
    
    await connection.send_message({
        "type": "stream_resumed",
        "stream_id": stream_id,
        "timestamp": datetime.utcnow().isoformat()
    })


async def handle_get_stream_status(connection: WebSocketConnection, data: Dict[str, Any]):
    """处理获取流状态消息"""
    stream_id = data.get("stream_id")
    
    if stream_id:
        if stream_id not in connection.stream_sessions:
            await connection.send_error(f"Stream {stream_id} not found", "STREAM_NOT_FOUND")
            return
        
        stream_manager = connection.stream_sessions[stream_id]
        status_data = stream_manager.get_stream_info()
        
        await connection.send_message({
            "type": "stream_status",
            "stream_id": stream_id,
            "status": status_data,
            "timestamp": datetime.utcnow().isoformat()
        })
    else:
        # 返回所有流的状态
        streams_status = {}
        for sid, stream_manager in connection.stream_sessions.items():
            streams_status[sid] = stream_manager.get_stream_info()
        
        await connection.send_message({
            "type": "all_streams_status",
            "streams": streams_status,
            "timestamp": datetime.utcnow().isoformat()
        })


async def handle_ping(connection: WebSocketConnection, data: Dict[str, Any]):
    """处理ping消息"""
    await connection.send_message({
        "type": "pong",
        "timestamp": datetime.utcnow().isoformat()
    })


async def setup_stream_callbacks(
    connection: WebSocketConnection,
    stream_manager: StreamManager,
    stream_id: str
):
    """设置流回调函数"""
    
    async def on_chunk_added(data, duration):
        await connection.send_stream_update(stream_id, "chunk_added", {
            "data_size": len(data),
            "duration": duration
        })
    
    async def on_buffer_update(buffer_info):
        await connection.send_stream_update(stream_id, "buffer_update", buffer_info)
    
    async def on_status_change(status):
        await connection.send_stream_update(stream_id, "status_change", {
            "status": status.value
        })
    
    async def on_error(error_message):
        await connection.send_stream_update(stream_id, "error", {
            "error_message": error_message
        })
    
    # 添加回调
    stream_manager.add_callback("on_chunk", on_chunk_added)
    stream_manager.add_callback("on_buffer_update", on_buffer_update)
    stream_manager.add_callback("on_status_change", on_status_change)
    stream_manager.add_callback("on_error", on_error)


async def heartbeat_loop(connection: WebSocketConnection):
    """心跳检测循环"""
    try:
        while True:
            await asyncio.sleep(30)  # 30秒心跳
            await connection.heartbeat()
    except asyncio.CancelledError:
        pass
    except Exception as e:
        print(f"心跳检测错误: {e}")


# API端点
@router.get("/connections")
@with_error_handling("get_websocket_connections")
@rate_limit(limit=10, window=60)
async def get_websocket_connections(
    current_user: Dict[str, Any] = Depends(get_current_user)
):
    """获取WebSocket连接信息"""
    try:
        user_id = current_user.get("sub")
        connections = websocket_manager.get_user_connections(user_id)
        
        connections_info = []
        for conn in connections:
            streams_info = {}
            for stream_id, stream_manager in conn.stream_sessions.items():
                streams_info[stream_id] = stream_manager.get_stream_info()
            
            connections_info.append({
                "connection_id": conn.connection_id,
                "connected_at": conn.connected_at.isoformat(),
                "last_activity": conn.last_activity.isoformat(),
                "active_streams": len(conn.stream_sessions),
                "streams": streams_info
            })
        
        return {
            "user_id": user_id,
            "total_connections": len(connections),
            "connections": connections_info
        }
        
    except Exception as e:
        raise BusinessLogicError(f"获取WebSocket连接信息失败: {str(e)}")


@router.post("/broadcast/{user_id}")
@with_error_handling("broadcast_to_user")
@require_permission("broadcast_message")
async def broadcast_to_user(
    user_id: str,
    message: Dict[str, Any],
    current_user: Dict[str, Any] = Depends(get_current_user)
):
    """向指定用户广播消息"""
    try:
        await websocket_manager.broadcast_to_user(user_id, {
            **message,
            "broadcast_from": current_user.get("sub"),
            "timestamp": datetime.utcnow().isoformat()
        })
        
        # 记录审计日志
        audit_logger.log_user_action(
            user_id=current_user.get("sub"),
            action="broadcast_message",
            resource="websocket",
            details={"target_user_id": user_id, "message_type": message.get("type")}
        )
        
        return {
            "success": True,
            "message": f"消息已广播给用户 {user_id}",
            "user_connections": len(websocket_manager.get_user_connections(user_id))
        }
        
    except Exception as e:
        raise BusinessLogicError(f"广播消息失败: {str(e)}")


# 定期清理不活跃连接
async def cleanup_websocket_connections():
    """定期清理不活跃的WebSocket连接"""
    while True:
        try:
            cleaned = await websocket_manager.cleanup_inactive_connections(timeout_seconds=300)
            if cleaned > 0:
                logging.info(f"清理了 {cleaned} 个不活跃的WebSocket连接")
            
            await asyncio.sleep(60)  # 每分钟检查一次
            
        except Exception as e:
            logging.error(f"清理WebSocket连接错误: {e}")
            await asyncio.sleep(60)


# 启动清理任务
asyncio.create_task(cleanup_websocket_connections())
